import math
from pathlib import Path

import numpy as np
import pandas as pd

from .utils import timer
from .typing import *


def range_kind(surface, volume):
    info_df = pd.DataFrame(list(surface.items()), columns=['Pair', 'Tuple'])
    info_df[['Surface', 'Distance', 'Type1', 'Type2']] = pd.DataFrame(info_df['Tuple'].tolist(), index=info_df.index)
    info_df[['Residue1', 'Atom1', 'Residue2', 'Atom2']] = info_df['Pair'].apply('_'.join).str.split('_', expand=True)
    
    # 把体积、原子类型、相互作用类型加上去
    info_df['Volume'] = 0
    info_df['Range']  = 0
    info_df['Kind']   = 'OTHER'
    info_df['Dist1']  = info_df['Residue1'].apply(lambda x: int(x[1:-1]))
    info_df['Dist2']  = info_df['Residue2'].apply(lambda x: int(x[1:-1]))
    info_df['Dist']   = info_df['Dist1'] - info_df['Dist2']
    
    # 匹配体积
    Volu = pd.DataFrame(list(volume.items()), columns=['Pair', 'Volume'])
    info_df = info_df.merge(Volu[['Pair', 'Volume']], on='Pair', how='left')
    info_df = info_df.rename(columns={'Volume_y': 'Volume'})
    
    # 匹配相互作用类型
    info_df['Kind'] = np.where((info_df['Type1'] == 'I') & # 氢键
                               (info_df['Type2'] == 'I') &
                               (1.5 <= info_df['Distance']) &
                               (info_df['Distance'] <= 3.5), 'HB', info_df['Kind'])
    info_df['Kind'] = np.where((info_df['Type1'] == 'I') & # 氢键
                               (info_df['Type2'] == 'II') &
                               (1.5 <= info_df['Distance']) &
                               (info_df['Distance'] <= 3.5), 'HB', info_df['Kind'])
    info_df['Kind'] = np.where((info_df['Type1'] == 'I') & # 氢键
                               (info_df['Type2'] == 'III') &
                               (1.5 <= info_df['Distance']) &
                               (info_df['Distance'] <= 3.5), 'HB', info_df['Kind'])
    info_df['Kind'] = np.where((info_df['Type1'] == 'II') & # 氢键
                               (info_df['Type2'] == 'I') &
                               (1.5 <= info_df['Distance']) &
                               (info_df['Distance'] <= 3.5), 'HB', info_df['Kind'])
    info_df['Kind'] = np.where((info_df['Type1'] == 'II') & # 氢键
                               (info_df['Type2'] == 'III') &
                               (1.5 <= info_df['Distance']) &
                               (info_df['Distance'] <= 3.5), 'HB', info_df['Kind'])
    info_df['Kind'] = np.where((info_df['Type1'] == 'III') & # 氢键
                               (info_df['Type2'] == 'I') &
                               (1.5 <= info_df['Distance']) &
                               (info_df['Distance'] <= 3.5), 'HB', info_df['Kind'])
    info_df['Kind'] = np.where((info_df['Type1'] == 'III') & # 氢键
                               (info_df['Type2'] == 'II') &
                               (1.5 <= info_df['Distance']) &
                               (info_df['Distance'] <= 3.5), 'HB', info_df['Kind'])
    
    info_df['Kind'] = np.where((info_df['Type1'] == 'V') & # 芳香族相互作用
                               (info_df['Type2'] == 'V'), 'AROM', info_df['Kind'])
    
    info_df['Kind'] = np.where((info_df['Type1'] == 'IV') & # 疏水相互作用
                               (info_df['Type2'] == 'IV'), 'PHOB', info_df['Kind'])
    info_df['Kind'] = np.where((info_df['Type1'] == 'IV') & # 疏水相互作用
                               (info_df['Type2'] == 'V'),  'PHOB', info_df['Kind'])
    info_df['Kind'] = np.where((info_df['Type1'] == 'V') &  # 疏水相互作用
                               (info_df['Type2'] == 'IV'), 'PHOB', info_df['Kind'])
    
    info_df['Kind'] = np.where((info_df['Type1'] == 'I') & # 不稳定相互作用
                               (info_df['Type2'] == 'IV'), 'DC', info_df['Kind'])
    info_df['Kind'] = np.where((info_df['Type1'] == 'II') & # 不稳定相互作用
                               (info_df['Type2'] == 'II'), 'DC', info_df['Kind'])
    info_df['Kind'] = np.where((info_df['Type1'] == 'II') & # 不稳定相互作用
                               (info_df['Type2'] == 'IV'), 'DC', info_df['Kind'])
    info_df['Kind'] = np.where((info_df['Type1'] == 'II') & # 不稳定相互作用
                               (info_df['Type2'] == 'VIII'), 'DC', info_df['Kind'])
    info_df['Kind'] = np.where((info_df['Type1'] == 'III') & # 不稳定相互作用
                               (info_df['Type2'] == 'III'), 'DC', info_df['Kind'])
    info_df['Kind'] = np.where((info_df['Type1'] == 'III') & # 不稳定相互作用
                               (info_df['Type2'] == 'IV'), 'DC', info_df['Kind'])
    info_df['Kind'] = np.where((info_df['Type1'] == 'III') & # 不稳定相互作用
                               (info_df['Type2'] == 'VII'), 'DC', info_df['Kind'])
    info_df['Kind'] = np.where((info_df['Type1'] == 'IV') & # 不稳定相互作用
                               (info_df['Type2'] == 'I'), 'DC', info_df['Kind'])
    info_df['Kind'] = np.where((info_df['Type1'] == 'IV') & # 不稳定相互作用
                               (info_df['Type2'] == 'II'), 'DC', info_df['Kind'])
    info_df['Kind'] = np.where((info_df['Type1'] == 'IV') & # 不稳定相互作用
                               (info_df['Type2'] == 'III'), 'DC', info_df['Kind'])
    info_df['Kind'] = np.where((info_df['Type1'] == 'VII') & # 不稳定相互作用
                               (info_df['Type2'] == 'III'), 'DC', info_df['Kind'])
    info_df['Kind'] = np.where((info_df['Type1'] == 'VII') & # 不稳定相互作用
                               (info_df['Type2'] == 'VII'), 'DC', info_df['Kind'])
    info_df['Kind'] = np.where((info_df['Type1'] == 'VIII') & # 不稳定相互作用
                               (info_df['Type2'] == 'II'), 'DC', info_df['Kind'])
    info_df['Kind'] = np.where((info_df['Type1'] == 'VIII') & # 不稳定相互作用
                               (info_df['Type2'] == 'VIII'), 'DC', info_df['Kind'])
    
    info_df['Kind'] = np.where((info_df['Atom1'] == 'C') & # 共价的肽键，和后一个残基
                               (info_df['Atom2'] == 'N') &
                               ((info_df['Residue1'].apply(lambda x: x[0]) ==
                                 info_df['Residue2'].apply(lambda x: x[0]))) &
                               (info_df['Dist'] == -1), 'PB', info_df['Kind'])
    info_df['Kind'] = np.where((info_df['Atom1'] == 'N') & # 共价的肽键，和前一个残基
                               (info_df['Atom2'] == 'C') &
                               ((info_df['Residue1'].apply(lambda x: x[0]) ==
                                 info_df['Residue2'].apply(lambda x: x[0]))) &
                               (info_df['Dist'] == 1), 'PB', info_df['Kind'])
    
    info_df['Kind'] = np.where((info_df['Atom1'] == 'SG') & # 共价的二硫键
                               (info_df['Atom2'] == 'SG') &
                               (1.95 <= info_df['Distance']) &
                               (info_df['Distance'] <= 2.1), 'SS', info_df['Kind'])
    
    # 匹配相互作用距离范围
    info_df['Range'] = np.where((abs(info_df['Dist']) <= 2) & # 短程
                                ((info_df['Residue1'].apply(lambda x: x[0]) ==
                                  info_df['Residue2'].apply(lambda x: x[0]))), 'S', info_df['Range'])
    info_df['Range'] = np.where((abs(info_df['Dist']) > 2) & # 中程
                                (abs(info_df['Dist']) <= 4) &
                                ((info_df['Residue1'].apply(lambda x: x[0]) ==
                                  info_df['Residue2'].apply(lambda x: x[0]))), 'M', info_df['Range'])
    info_df['Range'] = np.where((abs(info_df['Dist']) > 4) & # 长程
                                (abs(info_df['Dist']) <= 10) &
                                ((info_df['Residue1'].apply(lambda x: x[0]) ==
                                  info_df['Residue2'].apply(lambda x: x[0]))), 'L1', info_df['Range'])
    info_df['Range'] = np.where((abs(info_df['Dist']) > 10) &
                                (abs(info_df['Dist']) <= 20) &
                                ((info_df['Residue1'].apply(lambda x: x[0]) ==
                                  info_df['Residue2'].apply(lambda x: x[0]))), 'L2', info_df['Range'])
    info_df['Range'] = np.where((abs(info_df['Dist']) > 20) &
                                (abs(info_df['Dist']) <= 30) &
                                ((info_df['Residue1'].apply(lambda x: x[0]) ==
                                  info_df['Residue2'].apply(lambda x: x[0]))), 'L3', info_df['Range'])
    info_df['Range'] = np.where((abs(info_df['Dist']) > 30) &
                                (abs(info_df['Dist']) <= 40) &
                                ((info_df['Residue1'].apply(lambda x: x[0]) ==
                                  info_df['Residue2'].apply(lambda x: x[0]))), 'L4', info_df['Range'])
    info_df['Range'] = np.where((abs(info_df['Dist']) > 40) &
                                (abs(info_df['Dist']) <= 50) &
                                ((info_df['Residue1'].apply(lambda x: x[0]) ==
                                  info_df['Residue2'].apply(lambda x: x[0]))), 'L5', info_df['Range'])
    info_df['Range'] = np.where((abs(info_df['Dist']) > 50) &
                                ((info_df['Residue1'].apply(lambda x: x[0]) ==
                                  info_df['Residue2'].apply(lambda x: x[0]))), 'L6', info_df['Range'])
    # 链间
    info_df['Range'] = np.where((info_df['Residue1'].apply(lambda x: x[0]) !=
                                 info_df['Residue2'].apply(lambda x: x[0])), 'I', info_df['Range'])

    all_df = info_df[['Residue1', 'Atom1', 'Type1', 'Residue2', 'Atom2', 'Type2',
                      'Distance', 'Surface', 'Volume', 'Range', 'Kind']]

    return all_df


def summary(dihedral_angle, Residues):
    # 每个残基的综述
    EMPTY = np.zeros([len(Residues)])
    sum_df = pd.DataFrame()
    sum_df['Residue']   = list(Residues.keys())
    sum_df['Phi']       = EMPTY.copy()
    sum_df['Psi']       = EMPTY.copy()
    sum_df['Cova_Surf'] = EMPTY.copy()
    sum_df['Cova_Volu'] = EMPTY.copy()
    sum_df['NC_Surf']   = EMPTY.copy()
    sum_df['NC_Volu']   = EMPTY.copy()
    
    for i in range(len(sum_df)):
        a_res = sum_df.loc[i]
        # 加上二面角
        phi, psi = dihedral_angle[a_res['Residue']]
        # 转换为度数或设置为"None"
        sum_df.loc[i, 'Phi'] = '' if phi is None else f'{math.degrees(phi):.3f}'
        sum_df.loc[i, 'Psi'] = '' if psi is None else f'{math.degrees(psi):.3f}'
        
        # 非共价相互作用和共价相互作用的面积和体积
        interaction = Residues[a_res['Residue']]
        NC   = interaction[interaction['Kind'].isin(['HB', 'AROM', 'PHOB', 'DC', 'OTHER'])]
        Cova = interaction[interaction['Kind'].isin(['SS', 'PB'])]
        
        sum_df.loc[i, 'Cova_Surf'] = sum(Cova['Surface'])
        sum_df.loc[i, 'Cova_Volu'] = sum(Cova['Volume'])
        sum_df.loc[i, 'NC_Surf']   = sum(NC['Surface'])
        sum_df.loc[i, 'NC_Volu']   = sum(NC['Volume'])
    
    # 按照链和顺序排序
    sum_df['Chain']    = sum_df['Residue'].str[0]
    sum_df['Sequence'] = sum_df['Residue'].str[1:-1].astype(int)
    sorted_df = sum_df.sort_values(['Chain', 'Sequence'])
    res_df = sorted_df.drop(columns=['Chain', 'Sequence'])

    return res_df


def pdb_csv_sort(idx:int, name:str, dihedral_angle:Angles, surface:Surfaces, volume:Volumes, out_path:Path, disable_print=False):
    @timer(disable_print=disable_print)
    def make_result():
        out_dp = out_path / name
        out_dp.mkdir(parents=True, exist_ok=True)
        
        # 如果是只有一个构象
        if idx == -1:
            all_df = range_kind(surface, volume)
            all_df.to_csv(out_dp / f'_ALL_{name}.csv', index=0)
            
            Residues = dict(list(all_df.groupby('Residue1')))
            for res in Residues:
                Residues[res].to_csv(out_dp / f'{res}.csv', index=0)
            
            res_df = summary(dihedral_angle, Residues)
            res_df.to_csv(out_dp / f'_SUM_{name}.csv', index=0)

        # 如果有多个构象
        elif idx >= 0:
            model_dp = out_dp / f'{name}_{idx}'
            model_dp.mkdir(parents=True, exist_ok=True)
            
            all_df = range_kind(surface, volume)
            all_df.to_csv(model_dp / f'_ALL_{name}_{idx}.csv', index=0)
            
            Residues = dict(list(all_df.groupby('Residue1')))
            for res in Residues:
                Residues[res].to_csv(model_dp / f'{res}.csv', index=0)
              
            res_df = summary(dihedral_angle, Residues)
            res_df.to_csv(model_dp / f'_SUM_{name}_{idx}.csv', index=0)
    make_result()